#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于PyMySQL的增强数据库操作类
支持连接池、心跳检测、事务重试等机制，提高连接稳定性
"""

import pymysql
import json
import time
import threading
from typing import List, Dict, Any, Optional, Union
from queue import Queue, Empty
from threading import Lock, Timer

from 核心.配置文件 import settings


class ConnectionPool:
    """简单的数据库连接池"""
    
    def __init__(self, config: dict, max_connections: int = 5, min_connections: int = 2):
        self.config = config
        self.max_connections = max_connections
        self.min_connections = min_connections
        self.connections = Queue(maxsize=max_connections)
        self.active_connections = 0
        self.lock = Lock()
        self.heartbeat_timer = None
        
        # 初始化连接池
        self._init_pool()
        self._start_heartbeat()
    
    def _init_pool(self):
        """初始化连接池"""
        for _ in range(self.min_connections):
            try:
                conn = self._create_connection()
                if conn:
                    self.connections.put(conn)
                    self.active_connections += 1
            except Exception as e:
                print(f"初始化连接池失败: {e}")
    
    def _create_connection(self):
        """创建新的数据库连接"""
        try:
            conn = pymysql.connect(**self.config)
            return conn
        except Exception as e:
            print(f"创建数据库连接失败: {e}")
            return None
    
    def get_connection(self, timeout: int = 10):
        """获取数据库连接"""
        try:
            # 尝试从池中获取连接
            conn = self.connections.get(timeout=timeout)
            if conn and self._is_connection_alive(conn):
                return conn
            else:
                # 连接无效，关闭并创建新的
                if conn:
                    try:
                        conn.close()
                    except:
                        pass
                return self._create_connection()
        except Empty:
            # 池中没有可用连接，创建新的
            if self.active_connections < self.max_connections:
                with self.lock:
                    if self.active_connections < self.max_connections:
                        conn = self._create_connection()
                        if conn:
                            self.active_connections += 1
                            return conn
            # 达到最大连接数，等待
            return self.connections.get(timeout=timeout)
    
    def return_connection(self, conn):
        """归还连接到连接池"""
        if conn and self._is_connection_alive(conn):
            try:
                # 重置连接状态
                conn.rollback()
                self.connections.put(conn, timeout=1)
            except:
                # 如果归还失败，关闭连接
                try:
                    conn.close()
                except:
                    pass
                with self.lock:
                    self.active_connections -= 1
        else:
            # 连接无效，关闭并减少计数
            try:
                conn.close()
            except:
                pass
            with self.lock:
                self.active_connections -= 1
    
    def _is_connection_alive(self, conn):
        """检查连接是否有效"""
        try:
            if conn and conn.open:
                # 执行简单查询测试连接
                cursor = conn.cursor()
                cursor.execute("SELECT 1")
                cursor.fetchone()
                cursor.close()
                return True
        except:
            pass
        return False
    
    def _start_heartbeat(self):
        """启动心跳检测"""
        def heartbeat():
            while True:
                try:
                    # 检查所有连接
                    temp_connections = []
                    while not self.connections.empty():
                        try:
                            conn = self.connections.get_nowait()
                            if self._is_connection_alive(conn):
                                temp_connections.append(conn)
                            else:
                                try:
                                    conn.close()
                                except:
                                    pass
                                with self.lock:
                                    self.active_connections -= 1
                        except Empty:
                            break
                    
                    # 重新放回有效连接
                    for conn in temp_connections:
                        self.connections.put(conn)
                    
                    # 补充连接
                    while self.active_connections < self.min_connections:
                        conn = self._create_connection()
                        if conn:
                            self.connections.put(conn)
                            with self.lock:
                                self.active_connections += 1
                        else:
                            break
                    
                    time.sleep(30)  # 每30秒检查一次
                except Exception as e:
                    print(f"心跳检测异常: {e}")
                    time.sleep(30)
        
        heartbeat_thread = threading.Thread(target=heartbeat, daemon=True)
        heartbeat_thread.start()
    
    def close_all(self):
        """关闭所有连接"""
        while not self.connections.empty():
            try:
                conn = self.connections.get_nowait()
                try:
                    conn.close()
                except:
                    pass
            except Empty:
                break
        self.active_connections = 0


class MySQLCRUD:
    """增强的MySQL数据库操作类，支持连接池和自动重连"""
    
    def __init__(self, host: str, user: str, password: str, database: str, 
                 port: int = 3306, charset: str = 'utf8mb4', autocommit: bool = False,
                 max_retries: int = 3, connection_timeout: int = 60, 
                 pool_size: int = 5, **kwargs):
        """
        初始化数据库连接
        
        :param host: 数据库主机地址
        :param user: 用户名
        :param password: 密码
        :param database: 数据库名
        :param port: 端口号，默认3306
        :param charset: 字符集，默认utf8mb4
        :param autocommit: 是否自动提交，默认False
        :param max_retries: 最大重试次数，默认3
        :param connection_timeout: 连接超时时间，默认60秒
        :param pool_size: 连接池大小，默认5
        """
        self.config = settings.DB_CONFIG
        self.max_retries = max_retries
        self.pool_size = pool_size
        
        # 初始化连接池
        self.connection_pool = ConnectionPool(self.config, pool_size, max(2, pool_size // 2))
        
        # 单连接模式（兼容性）
        self.connection = None
        self.lock = Lock()
        
        print(f"✅ 数据库连接池初始化成功 (池大小: {pool_size})")
    
    def _get_connection(self):
        """获取数据库连接"""
        return self.connection_pool.get_connection()
    
    def _return_connection(self, conn):
        """归还数据库连接"""
        self.connection_pool.return_connection(conn)
    
    def is_connected(self) -> bool:
        """检查连接状态"""
        try:
            conn = self._get_connection()
            if conn:
                self._return_connection(conn)
                return True
        except:
            pass
        return False
    
    def execute_query(self, sql: str, params: Optional[Union[tuple, list, dict]] = None) -> List[Dict[str, Any]]:
        """
        执行查询SQL，返回字典列表（带重连机制）
        
        :param sql: SQL语句
        :param params: 参数
        :return: 查询结果列表
        """
        for attempt in range(self.max_retries):
            conn = None
            cursor = None
            try:
                conn = self._get_connection()
                if not conn:
                    print(f"无法获取数据库连接，重试 {attempt + 1}/{self.max_retries}")
                    time.sleep(2)
                    continue
                
                cursor = conn.cursor(pymysql.cursors.DictCursor)
                cursor.execute(sql, params)
                result = cursor.fetchall()
                cursor.close()
                self._return_connection(conn)
                return result

            except Exception as e:
                print(f"查询执行错误 (尝试 {attempt + 1}/{self.max_retries}): {e}")
                if cursor:
                    cursor.close()
                if conn:
                    self._return_connection(conn)
                    
                if attempt < self.max_retries - 1:
                    print("等待2秒后重试...")
                    time.sleep(2)
                else:
                    print("达到最大重试次数，查询失败")
                    return []

        return []
    
    def execute_update(self, sql: str, params: Optional[Union[tuple, list, dict]] = None) -> int:
        """
        执行更新SQL（INSERT/UPDATE/DELETE），返回影响行数（带重连机制）
        
        :param sql: SQL语句
        :param params: 参数
        :return: 影响的行数
        """
        for attempt in range(self.max_retries):
            conn = None
            cursor = None
            try:
                conn = self._get_connection()
                if not conn:
                    print(f"无法获取数据库连接，重试 {attempt + 1}/{self.max_retries}")
                    time.sleep(2)
                    continue
                
                cursor = conn.cursor()
                affected_rows = cursor.execute(sql, params)
                conn.commit()
                cursor.close()
                self._return_connection(conn)
                return affected_rows

            except Exception as e:
                print(f"更新执行错误 (尝试 {attempt + 1}/{self.max_retries}): {e}")
                if cursor:
                    cursor.close()
                if conn:
                    try:
                        conn.rollback()
                    except:
                        pass
                    self._return_connection(conn)
                    
                if attempt < self.max_retries - 1:
                    print("等待2秒后重试...")
                    time.sleep(2)
                else:
                    print("达到最大重试次数，更新失败")
                    return 0

        return 0
    
    def execute_batch(self, sql: str, params_list: List[Union[tuple, list, dict]]) -> int:
        """
        执行批量操作，返回影响行数（带重连机制和事务重试）
        
        :param sql: SQL语句
        :param params_list: 参数列表
        :return: 影响的行数
        """
        for attempt in range(self.max_retries):
            conn = None
            cursor = None
            try:
                conn = self._get_connection()
                if not conn:
                    print(f"无法获取数据库连接，重试 {attempt + 1}/{self.max_retries}")
                    time.sleep(2)
                    continue
                
                # 设置事务隔离级别，减少锁等待
                cursor = conn.cursor()
                cursor.execute("SET SESSION TRANSACTION ISOLATION LEVEL READ COMMITTED")
                cursor.execute("SET SESSION innodb_lock_wait_timeout = 30")
                
                # 分批执行，避免单次事务过大
                batch_size = 100
                total_affected = 0
                
                for i in range(0, len(params_list), batch_size):
                    batch = params_list[i:i + batch_size]
                    try:
                        affected_rows = cursor.executemany(sql, batch)
                        total_affected += affected_rows
                        conn.commit()  # 每批提交一次
                    except Exception as batch_error:
                        print(f"批次 {i//batch_size + 1} 执行失败: {batch_error}")
                        conn.rollback()
                        # 尝试逐条插入
                        for param in batch:
                            try:
                                cursor.execute(sql, param)
                                total_affected += 1
                            except Exception as single_error:
                                print(f"单条插入失败: {single_error}")
                        conn.commit()
                
                cursor.close()
                self._return_connection(conn)
                return total_affected

            except Exception as e:
                print(f"批量操作错误 (尝试 {attempt + 1}/{self.max_retries}): {e}")
                if cursor:
                    cursor.close()
                if conn:
                    try:
                        conn.rollback()
                    except:
                        pass
                    self._return_connection(conn)
                    
                if attempt < self.max_retries - 1:
                    print("等待2秒后重试...")
                    time.sleep(2)
                else:
                    print("达到最大重试次数，批量操作失败")
                    return 0

        return 0
    
    def execute_transaction(self, operations: List[tuple]) -> bool:
        """
        执行事务操作
        
        :param operations: 操作列表，每个元素是 (sql, params) 的元组
        :return: 是否成功
        """
        for attempt in range(self.max_retries):
            conn = None
            cursor = None
            try:
                conn = self._get_connection()
                if not conn:
                    print(f"无法获取数据库连接，重试 {attempt + 1}/{self.max_retries}")
                    time.sleep(2)
                    continue
                
                cursor = conn.cursor()
                
                # 开始事务
                cursor.execute("START TRANSACTION")
                
                # 执行所有操作
                for sql, params in operations:
                    cursor.execute(sql, params)
                
                # 提交事务
                conn.commit()
                cursor.close()
                self._return_connection(conn)
                return True

            except Exception as e:
                print(f"事务执行错误 (尝试 {attempt + 1}/{self.max_retries}): {e}")
                if cursor:
                    cursor.close()
                if conn:
                    try:
                        conn.rollback()
                    except:
                        pass
                    self._return_connection(conn)
                    
                if attempt < self.max_retries - 1:
                    print("等待2秒后重试...")
                    time.sleep(2)
                else:
                    print("达到最大重试次数，事务执行失败")
                    return False

        return False
    
    # 保持原有的便捷方法
    def insert(self, table: str, data: Dict[str, Any]) -> int:
        """插入单条数据"""
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
        values = tuple(data.values())
        
        return self.execute_update(sql, values)
    
    def insert_batch(self, table: str, data_list: List[Dict[str, Any]]) -> int:
        """批量插入数据"""
        if not data_list:
            return 0
        
        columns = list(data_list[0].keys())
        placeholders = ', '.join(['%s'] * len(columns))
        sql = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({placeholders})"
        
        values_list = []
        for data in data_list:
            values = [data.get(col) for col in columns]
            values_list.append(values)
        
        return self.execute_batch(sql, values_list)
    
    def upsert(self, table: str, data: Dict[str, Any], unique_keys: List[str]) -> int:
        """UPSERT操作（插入或更新）"""
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        
        update_columns = [col for col in data.keys() if col not in unique_keys]
        update_part = ', '.join([f"{col} = VALUES({col})" for col in update_columns])
        
        sql = f"INSERT INTO {table} ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {update_part}"
        values = tuple(data.values())
        
        return self.execute_update(sql, values)
    
    def upsert_batch(self, table: str, data_list: List[Dict[str, Any]], unique_keys: List[str]) -> int:
        """批量UPSERT操作"""
        if not data_list:
            return 0
        
        columns = list(data_list[0].keys())
        placeholders = ', '.join(['%s'] * len(columns))
        
        update_columns = [col for col in columns if col not in unique_keys]
        update_part = ', '.join([f"{col} = VALUES({col})" for col in update_columns])
        
        sql = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {update_part}"
        
        values_list = []
        for data in data_list:
            values = [data.get(col) for col in columns]
            values_list.append(values)
        
        return self.execute_batch(sql, values_list)
    
    def select(self, table: str, columns: str = "*", where: Optional[str] = None, 
               params: Optional[Union[tuple, list, dict]] = None, 
               order_by: Optional[str] = None, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """查询数据"""
        sql = f"SELECT {columns} FROM {table}"
        
        if where:
            sql += f" WHERE {where}"
        
        if order_by:
            sql += f" ORDER BY {order_by}"
        
        if limit:
            sql += f" LIMIT {limit}"
        
        return self.execute_query(sql, params)
    
    def select_one(self, table: str, columns: str = "*", where: Optional[str] = None, 
                   params: Optional[Union[tuple, list, dict]] = None) -> Optional[Dict[str, Any]]:
        """查询单条数据"""
        result = self.select(table, columns, where, params, limit=1)
        return result[0] if result else None
    
    def count(self, table: str, where: Optional[str] = None, 
              params: Optional[Union[tuple, list, dict]] = None) -> int:
        """统计记录数量"""
        sql = f"SELECT COUNT(*) as count FROM {table}"
        
        if where:
            sql += f" WHERE {where}"
        
        result = self.execute_query(sql, params)
        return result[0]['count'] if result else 0
    
    def update(self, table: str, data: Dict[str, Any], where: str, 
               params: Optional[Union[tuple, list, dict]] = None) -> int:
        """更新数据"""
        set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
        sql = f"UPDATE {table} SET {set_clause} WHERE {where}"
        
        values = tuple(data.values())
        if params:
            if isinstance(params, tuple):
                values += params
            else:
                values += (params,)
        
        return self.execute_update(sql, values)
    
    def delete(self, table: str, where: str, 
               params: Optional[Union[tuple, list, dict]] = None) -> int:
        """删除数据"""
        sql = f"DELETE FROM {table} WHERE {where}"
        return self.execute_update(sql, params)
    
    def table_exists(self, table: str) -> bool:
        """检查表是否存在"""
        sql = "SHOW TABLES LIKE %s"
        result = self.execute_query(sql, (table,))
        return len(result) > 0
    
    def get_table_info(self, table: str) -> List[Dict[str, Any]]:
        """获取表结构信息"""
        sql = f"DESCRIBE {table}"
        return self.execute_query(sql)
    
    def get_server_info(self) -> Optional[str]:
        """获取服务器信息"""
        try:
            conn = self._get_connection()
            if conn:
                info = conn.server_info
                self._return_connection(conn)
                return info
        except:
            pass
        return None
    
    def get_connection_id(self) -> Optional[int]:
        """获取连接ID"""
        try:
            conn = self._get_connection()
            if conn:
                cursor = conn.cursor()
                cursor.execute("SELECT CONNECTION_ID() as connection_id")
                result = cursor.fetchone()
                cursor.close()
                self._return_connection(conn)
                return result[0] if result else None
        except:
            pass
        return None
    
    def get_pool_status(self) -> str:
        """获取连接池状态"""
        try:
            if self.connection_pool:
                return f"连接池正常 (活跃连接: {self.connection_pool.active_connections})"
            else:
                return "连接池未初始化"
        except:
            return "连接池状态未知"
    
    def disconnect(self):
        """关闭数据库连接"""
        try:
            if self.connection_pool:
                self.connection_pool.close_all()
                print("数据库连接池已关闭")
        except Exception as e:
            print(f"关闭数据库连接池时出错: {e}")


# 使用示例
if __name__ == "__main__":
    # 数据库配置 - 请替换为您的实际数据库信息
    db_config = {
        'host': 'localhost',
        'database': 'test_db',
        'user': 'test_user',
        'password': 'test_password',
        'port': 3306,
        'charset': 'utf8mb4',
        'max_retries': 3,
        'connection_timeout': 60,
        'pool_size': 5
    }

    # 创建实例
    db = MySQLCRUD(**db_config)

    # 检查连接是否成功
    if not db.is_connected():
        print("数据库连接失败，程序退出")
        exit(1)

    # 显示连接状态
    print(f"连接状态: {db.get_pool_status()}")
    
    # 显示服务器信息
    server_info = db.get_server_info()
    if server_info:
        print(f"MySQL服务器版本: {server_info}")

    try:
        # 插入数据
        print("插入数据:")
        user_data = {'name': '张三', 'email': 'zhangsan@example.com', 'age': 25}
        db.insert('users', user_data)

        # 批量插入数据
        print("批量插入数据:")
        users_data = [
            {'name': '李四', 'email': 'lisi@example.com', 'age': 30},
            {'name': '王五', 'email': 'wangwu@example.com', 'age': 35}
        ]
        db.insert_batch('users', users_data)

        # 查询数据
        print("\n查询所有用户:")
        users = db.select('users')
        if users:
            for user in users:
                print(user)

    except Exception as e:
        print(f"程序执行错误: {e}")
    finally:
        # 关闭连接
        db.disconnect()